import time
import traceback
from uuid import uuid4
from loguru import logger

from fastapi import FastAPI, Request
from fastapi.responses import Response
from fastapi.middleware.cors import CORSMiddleware
from starlette.concurrency import iterate_in_threadpool

from core.response import resp_500
from core.logger import x_trace_id
from core.security import x_user_name, security
from common.database import SessionLocal
from app.audit.handlers.audit_api_record_handler import AuditApiRecordHandler


def register_process_time(app: FastAPI):
    @app.middleware("http")
    async def add_process_time_header(request: Request, call_next):
        start_time = time.time()
        response = await call_next(request)
        process_time = time.time() - start_time
        response.headers["X-Process-Time"] = str(process_time)
        return response


def register_trace_id(app: FastAPI):
    @app.middleware("http")
    async def add_trace_id(request: Request, call_next):
        trace_id = request.headers.get("x_trace_id", "")
        if trace_id == "":
            trace_id = uuid4().hex
        x_trace_id.set(str(trace_id))
        response = await call_next(request)
        response.headers["x_trace_id"] = str(trace_id)
        return response



def register_user_name(app: FastAPI):
    @app.middleware("http")
    async def add_user_name(request: Request, call_next):
        user_name = ""
        authorization = request.headers.get("Authorization", "")
        scheme, _, token = authorization.partition(" ")
        if scheme.lower() == "bearer" and token:
            payload = security.verify_access_token(token)
            if payload:
                user_name = payload.get("username","")
        x_user_name.set(str(user_name))
        response = await call_next(request)
        response.headers["x_user_name"] = str(user_name)
        return response



def register_request_record(app: FastAPI):
    async def set_body(request: Request):
        receive_ = await request._receive()
        async def receive():
            return receive_
        request._receive = receive
        
    async def get_request_text(request: Request):
        try:
            # 获取request body 
            content_type = request.headers.get("content-type", "")
            if 'multipart/form-data; boundary=' in content_type:
                request_text = "{'file_upload': 1}"
            else:
                await set_body(request)
                request_text = await request.body()
                request_text = str(request_text.decode("utf-8"))[:255]
                
            if request.url.path == "/ops/api/v1/login":
                # 去掉登陆的密码打印
                request_text = request_text.split("&")[0]
                
        except Exception as e:
            logger.error(f"get request text error=>{e}")
            return False
        return request_text
    
    async def get_response_text(response: Response):
        try:
            # 获取response body 
            response_body = [chunk async for chunk in response.body_iterator]
            response.body_iterator = iterate_in_threadpool(iter(response_body))
            response_text = (b''.join(response_body[:100])).decode()[:255]
        except Exception as e:
            logger.error(f"get response text error=>{e}")
            return False
        return response_text
    
    
    def save_record(request: Request, response: Response, request_text:str, response_text: str):
        try:
            # 记录到数据库审计
            data = {
                "user_name": x_user_name.get(),
                "user_ip": request.client.host,
                "user_agent": request.headers.get('user-agent',"")[:255],
                "path": request.url.path,
                "method": request.method,
                "url": request.url,
                "request_text": request_text,
                "response_text": response_text,
                "response_code": response.status_code,
                "x_trace_id": x_trace_id.get()
            }
                      
            # 保存记录信息
            audit_api_record_handler = AuditApiRecordHandler()
            db = SessionLocal()
            audit_api_record_handler.add_audit_api_record(db, data)
            db.close()
            
            # 打印记录信息
            logger.debug(f"request_record=>{data}")
        except Exception as e:
            logger.error(f"request_record error=>{e}")
            return False
        
        return True
        
        

    @app.middleware("http")
    async def request_record(request: Request, call_next):
        # 获取request体的文本 
        request_text = await get_request_text(request)
        
        # 执行接口处理函数
        response = await call_next(request)
        
        # 获取response体的文本
        response_text = await get_response_text(response)

        
        # 保存请求记录
        save_record(request, response, request_text, response_text)        
        
        return response
 


    # def register_exception(app: FastAPI):
    #     """
    #     全局异常捕获
    #     :param app:
    #     :return:
    #     """
    #
    #     # 捕获自定义异常
    #     @app.exception_handler(PostParamsError)
    #     async def query_params_exception_handler(request: Request, exc: PostParamsError):
    #         """
    #         捕获 自定义抛出的异常
    #         :param request:
    #         :param exc:
    #         :return:
    #         """
    #         logger.error(f"参数查询异常\nURL:{request.url}\nHeaders:{request.headers}\n{traceback.format_exc()}")
    #         return JSONResponse(
    #             status_code=status.HTTP_400_BAD_REQUEST,
    #             content={"code": 400, "data": {"tip": exc.err_desc}, "message": "fail"},
    #         )
    #
    #     @app.exception_handler(TokenAuthError)
    #     async def token_exception_handler(request: Request, exc: TokenAuthError):
    #         logger.error(f"参数查询异常\nURL:{request.url}\nHeaders:{request.headers}\n{traceback.format_exc()}")
    #         return JSONResponse(
    #             status_code=status.HTTP_400_BAD_REQUEST,
    #             content={"code": 400, "data": None, "message": exc.err_desc},
    #         )
    #
    #     # 捕获参数 验证错误
    #     @app.exception_handler(RequestValidationError)
    #     async def validation_exception_handler(request: Request, exc: RequestValidationError):
    #         """
    #         捕获请求参数 验证错误
    #         :param request:
    #         :param exc:
    #         :return:
    #         """
    #         logger.error(f"参数错误\nURL:{request.url}\nHeaders:{request.headers}\n{traceback.format_exc()}")
    #         return JSONResponse(
    #             status_code=status.HTTP_400_BAD_REQUEST,
    #             content=jsonable_encoder({"code": 400, "data": {"tip": exc.errors()}, "body": exc.body, "message": "fail"}),
    #         )
    #
    # 捕获全部异常
    # @app.exception_handler(Exception)
    # async def all_exception_handler(request: Request, exc: Exception):
    #     logger.error(f"全局异常\nURL:{request.url}\nHeaders:{request.headers}\n{traceback.format_exc()}")
    #     return resp_500()


def register_cors(app: FastAPI) -> None:
    """
    支持跨域

    貌似发现了一个bug
    https://github.com/tiangolo/fastapi/issues/133

    :param app:
    :return:
    """
    app.add_middleware(
        CORSMiddleware,
        # allow_origins=['http://localhost:8081'],  # 有效, 但是本地vue端口一直在变化, 接口给其他人用也不一定是这个端口
        # allow_origins=['*'],   # 无效 bug allow_origins=['http://localhost:8081']
        allow_origin_regex='https?://.*',  # 改成用正则就行了
        allow_credentials=True,
        allow_methods=["*"],
        allow_headers=["*"],
    )


def register_middleware(app: FastAPI):
    """
    请求响应拦截 hook

    https://fastapi.tiangolo.com/tutorial/middleware/
    :param app:
    :return:
    """

    @app.middleware("http")
    async def logger_request(request: Request, call_next):
        # https://stackoverflow.com/questions/60098005/fastapi-starlette-get-client-real-ip
        logger.info(f"访问记录:{request.method} url:{request.url}\nheaders:{request.headers.get('user-agent')}"
                    f"\nIP:{request.client.host}")
        response = await call_next(request)
        return response
